package com.db.router.context;

import com.db.router.dynamic.CustomConnection;

import javax.sql.DataSource;
import java.util.*;

/**
 * @author tianjiashu（mac）
 * @ClassName ConnectionContext.java
 * @Description 连接上下文，用来保持多数据事务中产生的Connection
 * @createTime 2023年07月12日 22:20:00
 */
public class ConnectionContext {
    private static final ThreadLocal<Map<DataSource,CustomConnection>> MULTI_TRAN_CONNECTION = new ThreadLocal<>();

    public static void setConnection(DataSource dataSource,CustomConnection connection){
        if(MULTI_TRAN_CONNECTION.get() == null){
            MULTI_TRAN_CONNECTION.set(new HashMap<>());
        }
        MULTI_TRAN_CONNECTION.get().put(dataSource,connection);
    }

    public static Collection<CustomConnection> getConnections(){
        return MULTI_TRAN_CONNECTION.get().values();
    }

    public static void clear(){
        MULTI_TRAN_CONNECTION.get().clear();
    }

    public static CustomConnection getConnection(DataSource dataSource){
        Map<DataSource, CustomConnection> map = MULTI_TRAN_CONNECTION.get();
        if(Objects.isNull(map)){
            return null;
        }
        return map.get(dataSource);
    }
}
